
#include "Stdafx.h"
#include "HeapProxy.h"
#include "ConsoleOutput.h"
#include "callstack.h"
#include "HookAPI.h"
#include "modulestree.h"

LPVOID	( WINAPI *IAT_HeapAlloc)( HANDLE hHeap,DWORD dwFlags,SIZE_T dwBytes ) = NULL;
BOOL	( WINAPI *IAT_HeapFree)( HANDLE hHeap,DWORD dwFlags,LPVOID lpMem ) = NULL;
LPVOID	( WINAPI *IAT_HeapReAlloc )(HANDLE hHeap,DWORD dwFlags,LPVOID lpMem, SIZE_T dwBytes ) = NULL;
//HANDLE (WINAPI *IAT_HeapCreate)(DWORD flOptions,SIZE_T dwInitialSize,SIZE_T dwMaximumSize ) = NULL;



__forceinline void heap_callstack( int add_bytes )
{
	int __c_hms = 0;
	HMODULE __vec_hms[MAX_CALLSTACK];
#if 1
	__try
	{
		int* p = NULL;
		*p = 0;
	}
	__except( excep_filter( GetExceptionInformation() ) )
	{
		//printf( "Some exception occures.\n" );
	}
#else

	dump_callstack( NULL );
#endif



	int hn = get_mod_stack_size();
	HMODULE* hm = get_mod_stack_base();

	__c_hms = 0;
	__vec_hms[0] = 0;

	for( int i = hn-1;i>=0;i-- )
	{		
		if( hm[i] == 0 )
			break;

		if( __c_hms>0 )
		{
			if( hm[i] == __vec_hms[__c_hms-1] )
				continue;
		}

		__vec_hms[__c_hms] = hm[i];
		__c_hms++;
	}

	int idx = find_mod_path( __vec_hms,__c_hms );

	if( idx>=0 )
	{
		add_mod_bytes( idx,add_bytes );
	}
}

__forceinline void HandleHeapReAlloc( HANDLE hHeap,DWORD dwFlags,LPVOID lpMem, SIZE_T dwBytes )
{
	DWORD dwBytes1 = HeapSize( hHeap,dwFlags,lpMem );
	heap_callstack( dwBytes-1 );
	//ConsoleOutput( "[ReAlloc]\t bytes=%d\n",dwBytes );
}
__forceinline void HandleHeapAlloc( HANDLE hHeap,DWORD dwFlags,SIZE_T dwBytes )
{
	//ConsoleOutput( "[Alloc]\t bytes=%d\n",dwBytes );
	heap_callstack( dwBytes );
}

__forceinline void HandleHeapFree( HANDLE hHeap,DWORD dwFlags,LPVOID lpMem )
{
	DWORD dwBytes1 = HeapSize( hHeap,dwFlags,lpMem );
	//ConsoleOutput( "[Free]\t ptr=0x%08x\n",lpMem );
	heap_callstack( -dwBytes1 );
}

//HANDLE WINAPI New_HeapCreate(DWORD flOptions,SIZE_T dwInitialSize,SIZE_T dwMaximumSize )
//{
//	return IAT_HeapCreate( flOptions,dwInitialSize,dwMaximumSize );
//}


LPVOID WINAPI New_HeapAlloc( HANDLE hHeap,DWORD dwFlags,SIZE_T dwBytes )
{
	HandleHeapAlloc( hHeap,dwFlags,dwBytes );
	return IAT_HeapAlloc( hHeap,dwFlags,dwBytes );
}

BOOL WINAPI New_HeapFree( HANDLE hHeap,DWORD dwFlags,LPVOID lpMem )
{
	HandleHeapFree( hHeap,dwFlags,lpMem );
	return IAT_HeapFree( hHeap,dwFlags,lpMem );
}

LPVOID WINAPI New_HeapReAlloc( HANDLE hHeap,DWORD dwFlags,LPVOID lpMem, SIZE_T dwBytes )
{	
	HandleHeapReAlloc( hHeap,dwFlags,lpMem,dwBytes );
	return IAT_HeapReAlloc( hHeap,dwFlags,lpMem,dwBytes );
}

void HookHeapAPIs( HMODULE instance )
{
	HMODULE hKernel = ::GetModuleHandleA("kernel32.dll");
	//if( !HookImportedFunc( instance,"kernel32.dll","HeapFree",(PROC)New_HeapFree,(PROC*)&IAT_HeapFree ) )
	{
		HOOKAPI api = {"HeapFree",(DWORD)GetProcAddress(hKernel,"HeapFree"),(PROC)New_HeapFree,0 };
		HookIAT( instance,"kernel32.dll",&api );
		if( !IAT_HeapFree )
			IAT_HeapFree = (BOOL (__stdcall *)(HANDLE,DWORD,LPVOID))api.old_proc;
	}

	//if( !HookImportedFunc( instance, "kernel32.dll","HeapAlloc",(PROC)New_HeapAlloc,(PROC*)&IAT_HeapAlloc ))
	{
		HOOKAPI api = {"HeapAlloc",(DWORD)GetProcAddress(hKernel,"HeapAlloc"),(PROC)New_HeapAlloc,0 };
		HookIAT( instance,"kernel32.dll",&api );
		if(!IAT_HeapAlloc)
			IAT_HeapAlloc = (LPVOID (__stdcall *)(HANDLE,DWORD,SIZE_T))api.old_proc;
	}

	//if( !HookImportedFunc( instance,"kernel32.dll","HeapReAlloc",(PROC)New_HeapReAlloc,(PROC*)&IAT_HeapReAlloc ))
	{
		HOOKAPI api = {"HeapReAlloc",(DWORD)GetProcAddress(hKernel,"HeapReAlloc"),(PROC)New_HeapReAlloc,0 };
		HookIAT( instance,"kernel32.dll",&api );
		if(!IAT_HeapReAlloc)
			IAT_HeapReAlloc = (LPVOID (__stdcall *)(HANDLE,DWORD,LPVOID,SIZE_T))api.old_proc;
	}
}

void UnHookHeapAPIs( HMODULE instance )
{
	static PROC p1,p2,p3;

	if( IAT_HeapFree )
	//if( !HookImportedFunc( instance,"kernel32.dll","HeapFree",(PROC)IAT_HeapFree,(PROC*)&p1 ) )
	{
		HOOKAPI api = {"HeapFree",(DWORD)IAT_HeapFree, (PROC)IAT_HeapFree,(PROC)p1 };
		HookIAT( instance,"kernel32.dll",&api );
	}
	if(IAT_HeapAlloc)
	//if( !HookImportedFunc( instance, "kernel32.dll","HeapAlloc",(PROC)IAT_HeapAlloc,(PROC*)&p1 ))
	{
		HOOKAPI api = {"HeapAlloc",(DWORD)IAT_HeapAlloc,(PROC)IAT_HeapAlloc,(PROC)p2 };
		HookIAT( instance,"kernel32.dll",&api );
	}
	if(IAT_HeapReAlloc)
	//if( !HookImportedFunc( instance,"kernel32.dll","HeapReAlloc",(PROC)IAT_HeapReAlloc,(PROC*)&p3 ))
	{
		HOOKAPI api = {"HeapReAlloc",(DWORD)IAT_HeapReAlloc,(PROC)IAT_HeapReAlloc,(PROC)p3 };
		HookIAT( instance,"kernel32.dll",&api );
	}
}